import numpy as np
from fealpy.quadrature.GaussLegendreQuadrature import GaussLegendreQuadrature





class Post():
    def __init__(self,PDE,sh,tspace):
        mesh = tspace.mesh
        bc = np.array([0.5,0.5])
        pp = mesh.bc_to_point(bc) #(NE,2)


        alpha = PDE.get_alpha(pp) #(NE,)
        index, = np.nonzero(np.abs(alpha-PDE.alpha)<1e-15) #(NE0,)

        
        self.mesh = mesh
        self.edge_basis = tspace.edge_basis

        self.PDE = PDE

        self.sh = sh[:][tspace.edge_to_dof()[index,:,1:]] #(NE0,ldof,2)

        self.n = mesh.edge_unit_normal(index=index) #(NE0,2)


        self.qf = GaussLegendreQuadrature(tspace.p+3)
        self.area = mesh.entity_measure('edge')[index]

        self.index = index

    def g(self,bc):
        index = self.index
        phi = self.edge_basis(bc)[...,index,:,1:,:] #(NQ,NE0,ldof,2,tdim)
        sh = self.sh #(NE0,ldof,2)

        shape = list(phi.shape)
        shape[-1] = 2

        phin = np.zeros(shape,dtype=np.float64) #(NQ,NE0,ldof,2,gdim)

        phin[...,0] = np.einsum('...ijlk,ik->...ijl',phi[...,[0,2]],self.n)
        phin[...,1] = np.einsum('...ijlk,ik->...ijl',phi[...,[2,1]],self.n)
       

        val = np.einsum('ijk,...ijkl->...il',sh,phin) #(NQ,NE0,2)

        return val #(NQ,NE0,2)



    def get_Q(self):
        index = self.index
        bcs, ws = self.qf.get_quadrature_points_and_weights()
        area = self.area
        pp = self.mesh.bc_to_point(bcs,index=index) #(NQ,NE0,2)

        W_Q = self.PDE.W_Q(pp) #(NQ,NE0,2)

        g = self.g(bcs) #(NQ,NE0,2)

        val = np.einsum('i,ijk,ijk,j->',ws,W_Q,g,area)/self.PDE.rho_0

        return val


    def get_M(self):
        index = self.index
        bcs, ws = self.qf.get_quadrature_points_and_weights()
        area = self.area
        pp = self.mesh.bc_to_point(bcs,index=index) #(NQ,NE0,2)

        W_M = self.PDE.W_M(pp) #(NQ,NE0,2)

        g = self.g(bcs) #(NQ,NE0,2)

        val = -np.einsum('i,ijk,ijk,j->',ws,W_M,g,area)/self.PDE.rho_0

        return val
